package com.distribute.config.distributeLock;

import com.distribute.domain.entity.DistributedLockEntity;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.mongodb.core.FindAndModifyOptions;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.Scheduled;

import javax.annotation.Resource;
import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

@Configuration
@Slf4j
@EnableScheduling
public class DistributedLockConfiguration {
    @Resource
    private MongoTemplate mongoTemplate;
//    private ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();

    private final ThreadPoolExecutor es = (ThreadPoolExecutor) Executors.newFixedThreadPool(10);

    public boolean acquireLock(String key, String token, long expiration, long timeout, ScheduledExecutorService scheduledExecutorService) {
        Query query = Query.query(Criteria.where("_id").is(key));
        final long currentTimeMillis = System.currentTimeMillis();
        Update update = new Update()
                .setOnInsert("_id", key)
                .setOnInsert("startAt", currentTimeMillis)
                .setOnInsert("expireAt", currentTimeMillis + expiration)
                .setOnInsert("token", token);
        FindAndModifyOptions options = new FindAndModifyOptions().upsert(true).returnNew(true);
        DistributedLockEntity distributedLock = mongoTemplate.findAndModify(query, update, options, DistributedLockEntity.class);
        if (distributedLock != null && isLocked(distributedLock, token)) {
            log.info("Try to acquire lock for key {} with token {}, Locked: {}, datetime: {}", key, token, true, LocalDateTime.now());
            startRefreshLockSchedule(distributedLock, expiration, timeout, scheduledExecutorService);
            return true;
        }
        return false;
    }

    public void releaseLock(String key, String token, ScheduledExecutorService scheduledExecutorService) {
//        es.shutdown();
        scheduledExecutorService.shutdown();
        Query query = Query.query(Criteria.where("_id").is(key).and("token").is(token));
        mongoTemplate.remove(query, DistributedLockEntity.class);
        log.info("Release lock successful and record for key {} with token {}, datetime: {}", key, token, LocalDateTime.now());
    }

    private void startRefreshLockSchedule(DistributedLockEntity distributedLock, long expiration, long timeout, ScheduledExecutorService scheduledExecutorService) {
        es.execute(new Runnable() {
            @Override
            public void run() {
                scheduledExecutorService.scheduleAtFixedRate(new Runnable() {
                    @Override
                    public void run() {
                        refreshLock(distributedLock, expiration, timeout, scheduledExecutorService);
                    }
                }, expiration/2, expiration/2, TimeUnit.MILLISECONDS);
            }
        });
    }

    public void refreshLock(DistributedLockEntity distributedLock, long expiration, long timeout, ScheduledExecutorService scheduledExecutorService) {
        final long currentTimeMillis = System.currentTimeMillis();
        String key = distributedLock.getId();
        String token = distributedLock.getToken();
        if (currentTimeMillis >= distributedLock.getStartAt() + timeout) {
            log.info("Holding time of the lock timeout, release lock and record for key {} with token {}, datetime: {}", key, token, LocalDateTime.now());
            releaseLock(key, token, scheduledExecutorService);
            return;
        }
        Query query = Query.query(Criteria.where("_id").is(key).and("token").is(token));
        Update update = Update.update("expireAt", currentTimeMillis + expiration);
        mongoTemplate.updateFirst(query, update, DistributedLockEntity.class);
        log.info("Refresh lock successful and record for key {} with token {}, datetime: {}, expire time: {}", key, token, LocalDateTime.now(),
                currentTimeMillis + expiration);
    }

    private boolean isLocked(DistributedLockEntity distributedLock, String token) {
        return distributedLock.getToken() != null && distributedLock.getToken().equals(token);
    }

    @Scheduled(fixedDelay = 60000)
    private void releaseExpiredLocks() {
        Query query = Query.query(Criteria.where("expireAt").lt(System.currentTimeMillis()));
        mongoTemplate.findAllAndRemove(query, DistributedLock.class);
//        log.info("Auto remove the expire records successful, datetime: {}", LocalDateTime.now());
    }
}
